package com.leadnews.controller.chatgpt;

import com.alibaba.fastjson.JSON;
import com.leadnews.controller.chatgpt.ai.Answer;
import com.leadnews.controller.chatgpt.ai.Choices;
import com.leadnews.dtos.ChatGpt;
import com.leadnews.mapper.ChatgptFromMapper;
import com.leadnews.mapper.ChatgptKeyMapper;
import com.leadnews.pojo.ChatgptFrom;
import com.leadnews.pojo.ChatgptKey;
import com.leadnews.pojo.User;
import com.leadnews.service.ChatgptFromService;
import com.leadnews.utils.TokenFilter;
import com.leadnews.vo.ResponseResult;
import org.apache.http.HttpStatus;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.conn.ssl.NoopHostnameVerifier;
import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
import org.apache.http.conn.ssl.TrustStrategy;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.ssl.SSLContexts;
import org.apache.http.util.EntityUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import javax.net.ssl.SSLContext;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.rmi.ServerException;
import java.security.KeyManagementException;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping
@CrossOrigin(origins = "*",maxAge = 3600)
public class ChatGptController {
    @Autowired
    TokenFilter tokenFilter;
    @Autowired
    ChatgptFromMapper chatgptFromService;
    @Autowired
    ChatgptKeyMapper chatgptKeyMapper;
    @PostMapping("/chatgpt")
    public ResponseResult chatgpt(@ModelAttribute ChatGpt chatgpt, HttpServletRequest request) throws IOException, NoSuchAlgorithmException, KeyStoreException, KeyManagementException{
        User user = tokenFilter.getToken(request);
        ResponseResult result = new ResponseResult();
        String accessKey=null;
        List<ChatgptFrom> from=null;
        if (user==null){
            from = chatgptFromService.findFrom(chatgpt.getFrom());
            if (from.size()==0){
                result.setCode(0);
                result.setErrorMessage("无权限访问");
                return result;
            }
            if (from.get(0).getCount()==0){
                result.setCode(0);
                result.setErrorMessage("无调用次数");
                return result;
            }
            accessKey=from.get(0).getAccessKey();
        }else{
            return null;
        }
        try (CloseableHttpClient httpClient = HttpClientBuilder.create()
                .setSSLSocketFactory(getSslConnectionSocketFactory())
                .build()) {
            ResponseResult submit = submit(httpClient, getHttpPost(accessKey), chatgpt.getQuestion());
            if (submit.getCode()==200){
                from.get(0).setCount(from.get(0).getCount()-1);
                int i = chatgptFromService.updateById(from.get(0));
                List<ChatgptKey> keyList = chatgptKeyMapper.findBy(accessKey);
                keyList.get(0).setCount(keyList.get(0).getCount()+1);
                int i1 = chatgptKeyMapper.updateById(keyList.get(0));
                if (i>0&&i1>0){
                    return submit;
                }
            }
            return null;
        }
    }
    private static ResponseResult submit(CloseableHttpClient httpClient, HttpPost post,String question) throws IOException {
            StringEntity stringEntity = new StringEntity(getRequestJson(question), getContentType());
            post.setEntity(stringEntity);
            CloseableHttpResponse response=null;
            try {
                response = httpClient.execute(post);
            } catch (SocketTimeoutException e) {

            } catch (SocketException e) {
                System.out.println("-- warning: Connection reset!");

            } catch (Exception e) {
                System.out.println("-- warning: Please try again!");
            }
        return printAnswer(response);

    }

    private static ResponseResult printAnswer(CloseableHttpResponse response) throws IOException {
        ResponseResult result = new ResponseResult();
        if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
            String responseJson = EntityUtils.toString(response.getEntity());
            Answer answer = JSON.parseObject(responseJson, Answer.class);
            StringBuilder answers = new StringBuilder();
            List<Choices> choices = answer.getChoices();
            for (Choices choice : choices) {
                answers.append(choice.getText());
            }
            result.setCode(200);
            result.setData(answers);
            System.out.println(answers);
        } else if (response.getStatusLine().getStatusCode() == 429) {
            result.setCode(10);
            result.setErrorMessage("-- warning: Too Many Requests!");
            System.out.println("-- warning: Too Many Requests!");
        } else if (response.getStatusLine().getStatusCode() == HttpStatus.SC_INTERNAL_SERVER_ERROR) {
            result.setCode(10);
            result.setErrorMessage("------ Server error, program terminated! ------");
            throw new ServerException("------ Server error, program terminated! ------");
        } else {
            result.setCode(10);
            result.setErrorMessage("-- warning: Error, please try again!");
            System.out.println("-- warning: Error, please try again!");
        }
        return result;
    }

    private static ContentType getContentType() {
        return ContentType.create("text/json", "UTF-8");
    }

    private static String getRequestJson(String question) {
        return "{\"model\": \"text-davinci-003\", \"prompt\": \"" + question + "\", \"temperature\": 0, \"max_tokens\": 1024}";
    }

    private static HttpPost getHttpPost(String accessKey) throws IOException {
        String openAiKey = accessKey;
        String connectTimeout = "600000";
        String connectionRequestTimeout = "600000";
        String socketTimeout = "600000";
        HttpPost post = new HttpPost("https://api.openai.com/v1/completions");
        post.addHeader("Content-Type", "application/json");
        post.addHeader("Authorization", "Bearer " + openAiKey);
        RequestConfig requestConfig = RequestConfig.custom()
                .setConnectTimeout(Integer.parseInt(connectTimeout)).setConnectionRequestTimeout(Integer.parseInt(connectionRequestTimeout))
                .setSocketTimeout(Integer.parseInt(socketTimeout)).build();
        post.setConfig(requestConfig);
        return post;
    }

    private static SSLConnectionSocketFactory getSslConnectionSocketFactory() throws NoSuchAlgorithmException, KeyStoreException, KeyManagementException {
        TrustStrategy acceptingTrustStrategy = (x509Certificates, s) -> true;
        SSLContext sslContext = SSLContexts.custom().loadTrustMaterial(null, acceptingTrustStrategy).build();
        return new SSLConnectionSocketFactory(sslContext, new NoopHostnameVerifier());
    }
}
